{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!./build_and_push.sh yolov5"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Testing your algorithm on your local machine\n",
    "\n",
    "When you're packaging your first algorithm to use with Amazon SageMaker, you probably want to test it yourself to make sure it's working correctly. We use the [SageMaker Python SDK](https://github.com/aws/sagemaker-python-sdk) to test both locally and on SageMaker. For more examples with the SageMaker Python SDK, see [Amazon SageMaker Examples](https://github.com/awslabs/amazon-sagemaker-examples/tree/master/sagemaker-python-sdk). In order to test our algorithm, we need our dataset."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## SageMaker Python SDK Local Training\n",
    "To represent our training, we use the Estimator class, which needs to be configured in five steps. \n",
    "1. IAM role - our AWS execution role\n",
    "2. train_instance_count - number of instances to use for training.\n",
    "3. train_instance_type - type of instance to use for training. For training locally, we specify `local` or `local_gpu`.\n",
    "4. image_name - our custom PyTorch Docker image we created.\n",
    "5. hyperparameters - hyperparameters we want to pass.\n",
    "\n",
    "Let's start with setting up our IAM role. We make use of a helper function within the Python SDK. This function throw an exception if run outside of a SageMaker notebook instance, as it gets metadata from the notebook instance. If running outside, you must provide an IAM role with proper access stated above in [Permissions](#Permissions)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker import get_execution_role\n",
    "\n",
    "role = get_execution_role()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fit, Deploy, Predict\n",
    "\n",
    "Now that the rest of our estimator is configured, we can call `fit()` with the path to our local CIFAR10 dataset prefixed with `file://`. This invokes our PyTorch container with 'train' and passes in our hyperparameters and other metadata as json files in /opt/ml/input/config within the container to our program entry point defined in the Dockerfile.\n",
    "\n",
    "After our training has succeeded, our training algorithm outputs our trained model within the /opt/ml/model directory, which is used to handle predictions.\n",
    "\n",
    "We can then call `deploy()` with an instance_count and instance_type, which is 1 and `local`. This invokes our PyTorch container with 'serve', which setups our container to handle prediction requests as defined [here](https://github.com/aws/sagemaker-pytorch-container/blob/master/src/sagemaker_pytorch_container/serving.py#L103). What is returned is a predictor, which is used to make inferences against our trained model.\n",
    "\n",
    "After our prediction, we can delete our endpoint.\n",
    "\n",
    "We recommend testing and training your training algorithm locally first, as it provides quicker iterations and better debuggability."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # Lets set up our SageMaker notebook instance for local mode.\n",
    "# !/bin/bash ./utils/setup.sh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import os\n",
    "# import subprocess\n",
    "\n",
    "# instance_type = 'local'\n",
    "\n",
    "# if subprocess.call('nvidia-smi') == 0:\n",
    "#     ## Set type to GPU if one is present\n",
    "#     instance_type = 'local_gpu'\n",
    "    \n",
    "# print(\"Instance type = \" + instance_type)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# from sagemaker.estimator import Estimator\n",
    "\n",
    "# hyperparameters = {}\n",
    "\n",
    "# estimator = Estimator(role=role,\n",
    "#                       train_instance_count=1,\n",
    "#                       train_instance_type=instance_type,\n",
    "#                       image_uri='yolov5:latest',\n",
    "#                       hyperparameters=hyperparameters)\n",
    "\n",
    "# estimator.fit({'cfg': 'file:///home/ec2-user/SageMaker/yolov5_sagemaker/data/cfg/', 'weights': 'file:///home/ec2-user/SageMaker/yolov5_sagemaker/data/weights/', 'images': 'file:///home/ec2-user/SageMaker/yolov5_sagemaker/data/images/', 'labels': 'file:///home/ec2-user/SageMaker/yolov5_sagemaker/data/labels/'})\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO\n",
    "#predictor = estimator.deploy(1, instance_type)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# predictor.delete_endpoint()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Part 2: Training and Hosting your Algorithm in Amazon SageMaker\n",
    "Once you have your container packaged, you can use it to train and serve models. Let's do that with the algorithm we made above.\n",
    "\n",
    "## Set up the environment\n",
    "Here we specify the bucket to use and the role that is used for working with SageMaker."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# S3 prefix\n",
    "prefix = 'DEMO-pytorch-yolov5'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create the session\n",
    "\n",
    "The session remembers our connection parameters to SageMaker. We use it to perform all of our SageMaker operations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sagemaker as sage\n",
    "\n",
    "sess = sage.Session()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Upload the data for training\n",
    "\n",
    "We will use the tools provided by the SageMaker Python SDK to upload the data to a default bucket."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "WORK_DIRECTORY = '/home/ec2-user/SageMaker/yolov5_sagemaker/data/'\n",
    "\n",
    "data_location = sess.upload_data(WORK_DIRECTORY, key_prefix=prefix)\n",
    "\n",
    "inputs = {'cfg': data_location+'/cfg', 'weights': data_location+'/weights', 'images': data_location+'/images', 'labels': data_location+'/labels'}\n",
    "\n",
    "print(inputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training on SageMaker\n",
    "Training a model on SageMaker with the Python SDK is done in a way that is similar to the way we trained it locally. This is done by changing our train_instance_type from `local` to one of our [supported EC2 instance types](https://aws.amazon.com/sagemaker/pricing/instance-types/).\n",
    "\n",
    "In addition, we must now specify the ECR image URL, which we just pushed above.\n",
    "\n",
    "Finally, our local training dataset has to be in Amazon S3 and the S3 URL to our dataset is passed into the `fit()` call.\n",
    "\n",
    "Let's first fetch our ECR image url that corresponds to the image we just built and pushed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import boto3\n",
    "\n",
    "client = boto3.client('sts')\n",
    "account = client.get_caller_identity()['Account']\n",
    "\n",
    "my_session = boto3.session.Session()\n",
    "region = my_session.region_name\n",
    "\n",
    "algorithm_name = 'yolov5'\n",
    "\n",
    "if region.startswith('cn'):\n",
    "    ecr_image = '{}.dkr.ecr.{}.amazonaws.com.cn/{}:latest'.format(account, region, algorithm_name)\n",
    "else:\n",
    "    ecr_image = '{}.dkr.ecr.{}.amazonaws.com/{}:latest'.format(account, region, algorithm_name)\n",
    "\n",
    "print(ecr_image)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.estimator import Estimator\n",
    "\n",
    "hyperparameters = {}\n",
    "\n",
    "instance_type = 'ml.g4dn.xlarge'\n",
    "\n",
    "estimator = Estimator(role=role,\n",
    "                      instance_count=1,\n",
    "                      instance_type=instance_type,\n",
    "                      image_uri=ecr_image,\n",
    "                      hyperparameters=hyperparameters)\n",
    "\n",
    "estimator.fit(inputs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "instance_type = 'ml.m5.xlarge'\n",
    "predictor = estimator.deploy(1, instance_type)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Optional cleanup\n",
    "When you're done with the endpoint, you should clean it up.\n",
    "\n",
    "All of the training jobs, models and endpoints we created can be viewed through the SageMaker console of your AWS account."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "predictor.delete_endpoint()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Reference\n",
    "- [How Amazon SageMaker interacts with your Docker container for training](https://docs.aws.amazon.com/sagemaker/latest/dg/your-algorithms-training-algo.html)\n",
    "- [How Amazon SageMaker interacts with your Docker container for inference](https://docs.aws.amazon.com/sagemaker/latest/dg/your-algorithms-inference-code.html)\n",
    "- [CIFAR-10 Dataset](https://www.cs.toronto.edu/~kriz/cifar.html)\n",
    "- [SageMaker Python SDK](https://github.com/aws/sagemaker-python-sdk)\n",
    "- [Dockerfile](https://docs.docker.com/engine/reference/builder/)\n",
    "- [scikit-bring-your-own](https://github.com/awslabs/amazon-sagemaker-examples/blob/master/advanced_functionality/scikit_bring_your_own/scikit_bring_your_own.ipynb)\n",
    "- [SageMaker PyTorch container](https://github.com/aws/sagemaker-pytorch-container)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "conda_python3",
   "language": "python",
   "name": "conda_python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
